'''
Created on 2011-10-25

@author: alexandre

'''


import numpy as np
import texTools
from pbTest import pbTest, meta_estimator
from freqTest import sgnTestPValueH0, wsrTestPValueH0
import re
squareBracketRe = re.compile("\[(.+)\]")

escape = texTools.escapeSpecialChar

def getLossLD( algoNameL, lossLL):
    lossLD = {}
    lossLLT = zip(*lossLL)# transpose list of list
    for algoName, lossL in zip(algoNameL, lossLLT):
        lossLD[algoName] = np.asarray(lossL)
    return lossLD
    
    
def unparseLossLDD(txt):
    lossLDD = {}
    dsName = None
    algoNameL = []
    lossLL = None
    warningL = []
    for i, line in enumerate(txt.splitlines()):
                
        line = line.strip()
        if len(line) == 0: continue # empty line 
        if line[0] == '#': continue # comment
        #                fileContent.replace("<!--replace-after-form-->", appendContent ,1)
        # new dataset
        if line[0] == '[':

            # stack the accumulated lossLL
            if lossLL is not None:
                if len(lossLL) == 0: raise Exception( 'You have no samples for the dataset %s, starting on line %d.'%(dsName,i) ) 
                lossLDD[dsName] = getLossLD( algoNameL, lossLL )
                lossLL = None
                
                
            # get dsName
            match = squareBracketRe.search(line)
            if match is None: raise Exception( 'Error reading dataset name on line %d, with content "%s"'%(i,line) )
            dsName = match.group(1)
            algoNameL = None
            continue
        
        # extract algo Name
        if algoNameL is None:
            algoNameL = [ name.strip() for name in line.split('\t') ]
            
            if len(algoNameL) == 0:
                raise Exception( 'Problem reading algo name on line %d, with content "%s"'%(i, line ) )
            
            lossLine = True
            try:                valL = [float(algoName) for algoName in algoNameL]
            except ValueError:  lossLine = False
            if lossLine:
                warningL.append( 'Expecting a line of algo names, but got all numbers. Is it a loss line ?. Line %d, with content "%s"'%(i,line))
            
            lossLL = []
            continue 
        
        # extract lossL
        if lossLL is not None:
            valL =[ float(val) for val in line.split('\t') ]
            if len(valL) != len(algoNameL):
                raise Exception( 'Got %d loss values for %d algo on line %d with content "%s" for dataset %s'%(len(valL), len(algoNameL), i, line, dsName ) )
            lossLL.append(valL)
            
            continue
    
    lossLDD[dsName] = getLossLD( algoNameL, lossLL ) 
    
    return lossLDD, warningL
    
def extractLossLL(lossLD):
    lossLL = []
    algoNameL = [] 
    for algoName, lossL in lossLD.items():
        lossLL.append(lossL)
        algoNameL.append(algoName)
    return algoNameL, lossLL
        
def parseLossLDD(lossLDD):
    txtL = []
    for dsName, lossLD in lossLDD.iteritems():
        txtL.append("[%s]" % dsName)
        algoNameL, lossLL = extractLossLL(lossLD)
        
        txtL.append('\t'.join(algoNameL))
        
        for lossL in zip(*lossLL):
            txtL.append('\t'.join([ str(loss) for loss in lossL ]))
    return '\n'.join(txtL)

class TestInfo:
    """
    Contains the test data for all algorithms on all datasests.
    """
    
    
    def __init__(self, lossLDD, dLossEstimator=meta_estimator):
        """
        lossLDD[dsKey][algoKey][testIdx] represents the loss of the classifier trained by [algoKey] on [dsKey] and evaluated on test example [testIdx].
        dLossEstimator is the function that returns the probability that classifier h has a better expected loss than classifier g, given dLossL,
        where, for [dsKey],  dLossL[testIdx] = lossLDD[dsKey][algoA][testIdx] - lossLDD[dsKey][algoB][testIdx], 
        h is the classifier produced by algoA  and g is the classifier produced by algoB.
        """
        self.lossLDD = lossLDD
        self.dLossEstimator = dLossEstimator
        self.algoKeyL = self._getAlgoKeyL()
        self.dsKeyL = lossLDD.keys()
        self._checkIntegrety()
    
    def _getAlgoKeyL(self):
        algoKeySet = set()
        for lossLD in self.lossLDD.values():
            algoKeySet.update(lossLD.keys())
        return list(algoKeySet)
    
    def getDiff(self, algoA, algoB, mean=True):
        diffD = {}
        for dsKey, lossLD in self.lossLDD.items():
            try:
                dL = np.asarray(lossLD[algoA]) - np.asarray(lossLD[algoB])
                if mean:  diffD[dsKey] = dL.mean()
                else:     diffD[dsKey] = dL
                      
                
            except KeyError: pass
        return diffD
    

    def getPairwisePoissonBinomial(self, algoKeyL=None):
        if algoKeyL is None : algoKeyL = self.algoKeyL
        algoProbD = {}
        individualProbD = {}
        for algoA in algoKeyL:
            for algoB in algoKeyL:
                diffLD = self.getDiff(algoA, algoB, mean=False)
                
                
                pD, pAvsB = pbTest(diffLD, self.dLossEstimator)
                algoProbD[(algoA, algoB)] = pAvsB
                individualProbD[ (algoA, algoB) ] = pD
    
        return algoProbD, individualProbD
    
    
    def getPairwiseFreqTest(self, algoKeyL=None, test=sgnTestPValueH0):
        if algoKeyL is None : algoKeyL = self.algoKeyL
        confD = {}
        for algoA in algoKeyL:
            for algoB in algoKeyL:
                diffD = self.getDiff(algoA, algoB, mean=True)
                pVal = test(diffD.values(), twoTailed=False)
#                if conf < 0: conf = None
                confD[(algoA, algoB)] = pVal
            
        return confD
    
    def getExpectedLoss(self):
        table = {}
        for dsKey, lossLD in self.lossLDD.items():
            for algoKey, lossL in lossLD.items():
                table[ dsKey, algoKey ] = np.mean(lossL)
        return table 
    
    def _checkIntegrety(self):
        for lossLD in self.lossLDD.values():
            lenL = np.array([ len(lossL) for  lossL in lossLD.values() ])
            assert (lenL[0] == lenL).all(), ', '.join( [str(l) for l in lenL ] ) 
    
    def __str__(self):
        strL = []
        for dsKey, lossLD in self.lossLDD.items():
            strL.append(dsKey)
            for algoKey, lossL in lossLD.items():
                strL.append('%s : %d' % (algoKey, len(lossL)))
        return '\n'.join(strL)
    
class BasicDsMask:
    def __init__(self, name, n):
        self.n = n
        self.name = name
        
    def __len__(self):
        return self.n
    
    def __str__(self):
        return self.name


class HpGridSpace:
    
    def __init__(self, hpLD):
        self.hpLD = hpLD
        
    def __tex__(self):
        texStr = "\\begin{description}\n"
        for key, hpL in self.hpLD.items():
            if isinstance(hpL[0], float):
                hpStr = ', '.join([ '%.2g' % hp for hp in hpL ])
            else:
                hpStr = ', '.join([ '%s' % hp for hp in hpL ])
            texStr += "\\item[%s] $\in \{$ %s $\}$\n" % (escape(key), escape(hpStr))   
        texStr += "\\end{description}\n"
        return texStr
        
        

class DsInfo:

    def __init__(self, dsKey, dsName, trnSz, tstSz, nFeatures, comment=None):
        self.key = dsKey # used to make links between object
        self.name = dsName # the actual displayed name
        self.trnSz = trnSz
        self.tstSz = tstSz
        self.nFeatures = nFeatures
        self.comment = comment


class AlgoInfo:
    
    def __init__(self, key, name, description=None, hpSpace=None, metricD=None, metricCaptionD={}, dsCaptionD={}):
        self.name = name
        self.key = key
        self.hpSpace = hpSpace
        self.metricD = metricD
        self.metricCaptionD = metricCaptionD
        self.description = description 
        self.dsCaptionD = dsCaptionD

    def __tex__(self):
        texStrL = []
        
        if self.description is not None:
            texStrL.append("""\\paragraph{Description} %s""" % (escape(self.description)))
            
        if self.hpSpace is not None:
            texStrL.append("""\\paragraph{Hyperparameter space} \n %s""" % (self.hpSpace.__tex__()))
        
        if self.metricD is not None:
            metricTable = texTools.TableMap(self.metricD)
            metricTable.formatInfo(rowCaption=self.dsCaptionD.get, colCaption=self.metricCaptionD.get)
            texStrL.append(metricTable.__tex__())
        
        return '\n\n'.join(texStrL)
        
        

class Context:
    
    def __init__(self, name, dsInfoL, description=None):
        self.name = name
        self.dsInfoL = dsInfoL
        self.description = description
        self.captionD = {}
        for dsInfo in dsInfoL:
            self.captionD[dsInfo.key] = dsInfo.name
    
    def __tex__(self):
        pass




class Report:
    """
    Will produce a latex report for comparing learning algorithms out of testing data. 
    """
    
    def __init__(self, testInfo, context=None, algoInfoL=None,
                 showIndividualRisk=True, pbTest=True, wsrTest=False, sgnTest=False):
        
        """
        testInfo : parameter of type TestInfo to provide the testing information for all algorithms on all datasets.
        context : parameter of type Context to provide information about the list of datasets (optional)
        showIndividualRisk : Whether or not the riskTable should be shown (default True). 
        pbTest : Whether or not the Poisson binomial test should be shown (default True).
        wsrTest : Whether or not Wilcoxon signed rank test should be shown (default False). This test can produce unreliable results in some cases. 
        sgnTest : Whether or not the sign test should be shown (default False). This test is less powerful than the Poisson binomial test.
        """
        
        self.testInfo = testInfo
        self.algoInfoL = algoInfoL
        self.context = context
        self.pbTest = pbTest
        self.wsrTest = wsrTest
        self.sgnTest = sgnTest
        self.showIndividualRisk = showIndividualRisk
        
        self._buildReport()
        
    
    def _buildReport(self):
        
        if self.algoInfoL is not None:
            self.algoNameD = dict([ (algo.key, algo.name) for algo in self.algoInfoL ])
        else: 
            self.algoNameD = dict([ (key, key) for key in self.testInfo.algoKeyL ])
        
        
        self.algoNameTexD = dict( [ (key, escape(name)) for key, name in self.algoNameD.items() ] )
        self.dsNameTexD   = dict( [ (key, escape(key)) for key in self.testInfo.dsKeyL ] )
        print self.dsNameTexD
        
        self.texDoc = texTools.TexDoc()
        
        pbProbD, individualProbD = self.testInfo.getPairwisePoissonBinomial()
        self.algoKeyL = self._getAlgoOrder(pbProbD)
        
        
        if self.pbTest:
            pbProbTable = texTools.TableMap(pbProbD, self.algoKeyL, self.algoKeyL)
            pbProbTable.formatInfo(rowCaption=self.algoNameTexD.get, colCaption=self.algoNameTexD.get)
            pbProbTable.format = dict(zip(self.algoKeyL, ['%.2f']*len(self.algoKeyL)))
            pbProbTable.title = 'Pairwise Poisson binomial test, representing the probability that row is better than column.'
            self.texDoc.add(pbProbTable)

        if self.sgnTest:
            sgnTestTable = self._getFreqTable(sgnTestPValueH0)
            sgnTestTable.title = 'Pairwise sign test, representing $\Pr(H_0)$, where $\Pr(H_0)$ is the one tail $p$-value for $H_0$.'
            self.texDoc.add(sgnTestTable)

        if self.wsrTest:
            wsrTestTable = self._getFreqTable(wsrTestPValueH0)
            wsrTestTable.title = 'Pairwise Wilcoxon signed rank test, representing $ \Pr(H_0)$, where $\Pr(H_0)$ is the one tail $p$-value for $H_0$.'
            self.texDoc.add(wsrTestTable)
            
        if self.showIndividualRisk:
            lossD = self.testInfo.getExpectedLoss()
            riskTable = texTools.TableMap(lossD, colL=self.algoKeyL)
            riskTable.formatInfo(rowCaption=self.dsNameTexD.get, colCaption=self.algoNameTexD.get)
            riskTable.format = dict(zip(self.algoKeyL, ['%.3f']*len(self.algoKeyL)))
            riskTable.title = "Individual expected loss for each algorithm on each dataset."
            self._interlaceCompareInfo(riskTable, individualProbD, True)
            self.texDoc.add(riskTable)
        
    def append( self, tex_item):
        self.texDoc.add( tex_item )

    def _getFreqTable(self, test):
        algoNameD = self.algoNameTexD
        pValD = self.testInfo.getPairwiseFreqTest(test=test)
        table = texTools.TableMap(pValD, self.algoKeyL, self.algoKeyL)
        table.formatInfo(rowCaption=algoNameD.get, colCaption=algoNameD.get)
        table.format = dict(zip(algoNameD.keys(), ['%.2f']*len(algoNameD))) 
        return table

    def _getAlgoOrder(self, pBetterD):
        pSumD = {}
        for keyPair, p in pBetterD.items():
            i, _j = keyPair
            try :             pSumD[i] += p
            except KeyError : pSumD[i] = p
        

        sortedValuesL = [ (pSum, key) for key, pSum in pSumD.iteritems() ]
        sortedValuesL.sort(reverse=True)
        return [ key for pSum, key in sortedValuesL ]


    def _interlaceCompareInfo(self, riskTable, pBetterDD, barProb=True):
        algoKeyL = riskTable.colKeys()
        dsKeyL = riskTable.rowKeys()
        
        if barProb:
            self.texDoc.usepackage("tikz")
            self.texDoc.add("""\n\\providecommand{\\barProb}[1]{ \\tikz{ \\fill[color=black!30] (2*#1em,0.3ex) rectangle ( 2em, 1ex); \\fill[color=blue!60!black!80] (0em,0.3ex) rectangle ( 2*#1em, 1ex); } }\n""")
        
        compareD = {}
        colL = []
        i = None
        for j in algoKeyL:
            
            if i is None:  i = j ; continue
            
            dsD = {}
            for dsKey in dsKeyL:  
                try:
                    if barProb:  dsD[dsKey] = "\\barProb{%.2f}" % (pBetterDD[i, j][dsKey])
                    else:        dsD[dsKey] = "\\scriptsize{%.2f}" % (pBetterDD[i, j][dsKey])
                except KeyError : pass # missing values
            
            compareD[(i, j)] = dsD
            colL.append((i, j))
            i = j
        
        # formatting
        fmt = {}
        colCaption = self.algoNameTexD
        for compKey, dsD in compareD.iteritems():
            riskTable.insertCol(compKey, dsD, compKey[1])
            colCaption[compKey] = '$\succ$'
#            fmt[compKey] = '%.2f'
            
        
        riskTable.colCaption = colCaption.get
        riskTable.format.update(fmt)


# ------------------------------------
# test
# ------------------------------------


if __name__ == "__main__":
    from util import readFile
    
    testInfo = TestInfo(unparseLossLDD( readFile( 'lossLDD.txt' ) ))
    report = Report(testInfo, sgnTest=True, wsrTest=True, showIndividualRisk=True)
    report.texDoc.pdf()

    
    
#from hpExplore import evalUtil
#from hpExplore.viz.tableView import str2tupleMetric, regHp, extractMetric
#
#
#def loadTestData(testPath, argMinMetric='xv.mean.val.risk'):
#    
#    evalMap = evalUtil.loadEvalMapRec(testPath)
#    evalMap.argMin(argMinMetric, regFun=regHp)
#    lossBinLD, _hpTagD = extractMetric(evalMap, str2tupleMetric('full.tst.errStream'))
#    lossLD = {}
#    for key, lossBinL in lossBinLD.items():
#        lossLD[key] = np.unpackbits(lossBinL).astype(np.float)
#    return lossLD
#
#
#
#def buildTestInfo(testL):
#    lossLDD = {}
#    for testPath, algoKey in testL:
#        for dsKey, lossL in loadTestData(testPath).items():
#            lossLDD.setdefault(dsKey, {})[algoKey] = lossL
#    return TestInfo(lossLDD)
#
#if __name__ == "__main__":
#    
#    from fileUtil import readFile
#    
#    lossLDD = unparseLossLDD(readFile('lossLDD.txt')) 
#
#    
#    from os import path
#    from fileUtil import writeFile, writePklz
#    testFolder = path.expandvars('$HOME/data/testsImportants')
#    
#    svmTest = (path.join(testFolder, 'libsvm-RBF-AvgDist->pagerContext(split=0.50)[04Oct2010-20.52.30]'), 'svm')
#    parzenTest = (path.join(testFolder, 'parzen-RBF->pagerContext(split=0.50)[05Oct2010-13.32.09]'), 'parzen')
#    adaBoostTest = (path.join(testFolder, 'gkAdaBoost-stump->pagerContext(split=0.50)[12Oct2010-13.02.33]'), 'adaBoost')
##    blindTest    = ( path.join( testFolder, 'blindLearner->pagerContext(split=0.50)[12Oct2010-18.51.37]' ),  'blind' )
##    l2BoostTest  = ( path.join( testFolder, 'gkL2Boost-stump->pagerContext(split=0.50)[12Oct2010-21.23.29]' ),   'L2Boost' )
#    annTest = (path.join(testFolder, 'ann2h->pagerContext(split=0.50)[16Oct2010-11.27.41]'), 'ann')
#    
#    
#    testL = [parzenTest, adaBoostTest, svmTest, annTest]
#    testInfo = buildTestInfo(testL)
#    lossLDD = unparseLossLDD( readFile( 'lossLDD.txt' ) )
#    testInfo = TestInfo(lossLDD)
#    lossLDD = testInfo.lossLDD
#    
##    writePklz(testInfo.lossLDD, 'lossLDD.pklz')
##    txtLossLDD = parseLossLDD(lossLDD)
##    lossLDD = unparseLossLDD( txtLossLDD )
##
##    for dsName, lossLD in lossLDD.items():
##        for algoName, lossL in lossLD.items():
##            print  all(lossL == lossLDD_[dsName][algoName]), algoName, dsName
##    
##    writeFile(txtLossLDD, "lossLDD.txt")
#    
#    report = Report(testInfo, sgnTest=True, wsrTest=True, showIndividualRisk=True)
#
#
#    report.texDoc.pdf()
#
#    
